#include "Tools.h"

const wchar_t wsKeyValue[] = L"System\\ControlSet001\\Services\\FileHidingFilter\\Rules";
const wchar_t wsKeyName[] = L"HideFolderAction";


void PrintError(DWORD dwError)
{
    LPWSTR lpMsgBuf;

    FormatMessage(  FORMAT_MESSAGE_ALLOCATE_BUFFER | 
        FORMAT_MESSAGE_FROM_SYSTEM |
        FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        dwError,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPTSTR) &lpMsgBuf,
        0, NULL );

    std::wstring wsMessage(lpMsgBuf);
    std::wcerr<<wsMessage;
    LocalFree(lpMsgBuf);
}

bool CreateFolder(const wchar_t *str)
{
    if(!CreateDirectory(str, NULL))
    {
        DWORD dwError = GetLastError();
        PrintError(dwError);

        return false;
    }

    return true;
}

bool FindFile(const wchar_t *str)
{
    std::wstring wsParentFolder(str);
    std::size_t pos = wsParentFolder.rfind(L"\\");
    int nLength = wsParentFolder.size();

    if(pos == std::wstring::npos)
    {
        std::wcerr<<"You entered wrong path\n";
        return false;
    }

    if(pos == static_cast<size_t>(nLength-1))
    {
        std::wcerr<<"Enter path in such format: Path_to_file\\File_to_search\n";
        return false;
    }

    std::wstring wsFileName(wsParentFolder,pos+1,nLength-pos);

    HANDLE hFile;
    WIN32_FIND_DATA FileInformation;

    hFile = FindFirstFileW(str, &FileInformation);

    if(hFile == INVALID_HANDLE_VALUE)
    {
        DWORD dwError = GetLastError();

        switch(dwError)
        {
        case ERROR_FILE_NOT_FOUND:
            {
                std::wcerr<<"File not found\n";

                break;
            }
        case ERROR_PATH_NOT_FOUND:
            {
                std::wcerr<<"Path not found\n";

                break;
            }
        default:
            {
                PrintError(dwError);

                break;
            }

        }

        return false;
    }

    do
    {
        if(FileInformation.dwFileAttributes)
        {
            if(FileInformation.cFileName == wsFileName)
            {
                FindClose(hFile);

                return true;
            }
        }

    }while(FindNextFileW(hFile, &FileInformation) == TRUE);

    return false;

}

bool WriteFolderToRules(const wchar_t *str)
{
    CMultiSZObserver observer;

    try
    {
        EnumMultiSZValue(wsKeyValue, wsKeyName, &observer);
    }
    catch (std::runtime_error &e) 
    {
        std::string errorStr = e.what();

        if(errorStr == "Error key opening!" || errorStr == "Error value reading!")
        {
            HKEY hKeyMain = NULL;
            LONG lResult = RegCreateKeyW(HKEY_LOCAL_MACHINE, wsKeyValue, &hKeyMain);

            if(lResult != ERROR_SUCCESS)
            {
                std::wcerr<<"Error can't open key";

                return false;
            }

            KeyGuard keyGuard(hKeyMain);

        }
        else if(errorStr != "Error value reading!")
        {
            std::wcerr << "Runtime error: " << e.what( ) << std::endl;

            return false;
        }    
    }

    observer.OnMultiSZItemFound(str);

    try
    {
        observer.DeleteEmptyStrings();

        WriteMultiSZValue(wsKeyValue, wsKeyName, observer.GetVector());
    }
    catch (std::runtime_error &e)
    {
        std::wcerr << "Runtime error: " << e.what( ) << std::endl;

        return false;
    }

    return true;
}

bool ReloadRules(void)
{
    HANDLE hDriverHandle = CreateFileW( FileHidingFilterCDONameForUser,
        GENERIC_READ | GENERIC_WRITE,
        0,
        NULL,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL,
        NULL);

    if (INVALID_HANDLE_VALUE == hDriverHandle)
    {
        std::wcerr<<"Error can't attach to driver\n";

        return false;
    }

    DWORD dwBytesReturned = 0;

    if (!DeviceIoControl(hDriverHandle,
        IOCTL_FHF_RELOAD_RULES,
        NULL,
        0,
        NULL,
        0,
        &dwBytesReturned,
        NULL))
    {
        DWORD dwError = GetLastError();

        PrintError(dwError);

        CloseHandle(hDriverHandle);

        return false;
    }

    CloseHandle(hDriverHandle);

    return true;    
}